import pandas as pd
import numpy as np
import lightgbm as lgb
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import LabelEncoder # Although not used in final version, good to keep if needed
from sklearn.metrics import roc_auc_score
import gc # Garbage Collector for memory management
import re # For regular expressions in employmentLength cleaning

# --- 1. 数据加载 ---
print("--- 1. 数据加载 ---")
# 假设您已经将下载的训练集和测试集文件重命名为 train.csv 和 testA.csv
# 请根据您的实际文件名调整路径
try:
    train_df = pd.read_csv('train.csv') # 对应 s***********t.csv
    test_df = pd.read_csv('testA.csv')  # 对应 t***A.csv (测试集A)
    # 如果有测试集B，处理方式类似，但通常提交只针对一个测试集
    # testB_df = pd.read_csv('testB.csv') # 对应 t***n.csv
except FileNotFoundError:
    print("错误：请确保 'train.csv' 和 'testA.csv' 文件存在于当前目录下。")
    print("您可以将下载的文件重命名为 train.csv 和 testA.csv。")
    exit()

print(f"训练集大小: {train_df.shape}")
print(f"测试集大小: {test_df.shape}")

# 记录测试集的id，用于提交
test_id = test_df['id']

# 合并训练集和测试集，方便统一进行特征工程和预处理
# 注意：在合并前，先将训练集的标签列 'isDefault' 分离出来
y_train = train_df['isDefault']
train_df = train_df.drop(['isDefault', 'id'], axis=1) # 训练集去除id和标签
test_df = test_df.drop('id', axis=1) # 测试集去除id

# 确保训练集和测试集的列名一致，且顺序一致
# (在特征工程后会再次检查)
common_cols = list(set(train_df.columns) & set(test_df.columns))
all_data = pd.concat([train_df[common_cols], test_df[common_cols]], axis=0, ignore_index=True)
print(f"合并后的数据大小: {all_data.shape}")

del train_df, test_df
gc.collect()

# --- 2. 特征工程 ---
print("\n--- 2. 特征工程 ---")

# 2.1 日期特征处理
# issueDate (贷款发放的月份)
all_data['issueDate'] = pd.to_datetime(all_data['issueDate'], format='%Y-%m-%d')
all_data['issueDate_year'] = all_data['issueDate'].dt.year
all_data['issueDate_month'] = all_data['issueDate'].dt.month
all_data['issueDate_day'] = all_data['issueDate'].dt.day
all_data['issueDate_weekday'] = all_data['issueDate'].dt.weekday
all_data['issueDate_dayofyear'] = all_data['issueDate'].dt.dayofyear
all_data['issueDate_weekofyear'] = all_data['issueDate'].dt.isocalendar().week.astype(int)
# 贷款发放至今的月数 (可以根据一个固定日期计算，或者直接用年份差)
# 假设以2020年1月1日为基准
# all_data['issueDate_diff_months'] = (2020 - all_data['issueDate_year']) * 12 + (1 - all_data['issueDate_month'])
all_data = all_data.drop('issueDate', axis=1) # 删除原始日期列

# earliesCreditLine (借款人最早报告的信用额度开立的月份)
# 格式为 'YYYY-MM' 或 'Mon-YYYY'，需要统一处理
def parse_earliest_credit_line(date_str):
    if pd.isna(date_str):
        return np.nan
    try:
        # 尝试 'Mon-YYYY' 格式
        return pd.to_datetime(date_str, format='%b-%Y')
    except ValueError:
        try:
            # 尝试 'YYYY-MM' 格式
            return pd.to_datetime(date_str, format='%Y-%m')
        except ValueError:
            return np.nan # 无法解析的日期
            
all_data['earliesCreditLine'] = all_data['earliesCreditLine'].apply(parse_earliest_credit_line)

# 计算信用历史长度 (以月为单位)
# 假设以最近的issueDate为当前时间点，或者一个固定时间点
# 这里我们用 issueDate_year - earliesCreditLine_year 来近似
all_data['earliesCreditLine_year'] = all_data['earliesCreditLine'].dt.year
all_data['earliesCreditLine_month'] = all_data['earliesCreditLine'].dt.month
all_data['credit_history_length'] = (all_data['issueDate_year'] - all_data['earliesCreditLine_year']) * 12 + \
                                    (all_data['issueDate_month'] - all_data['earliesCreditLine_month'])
all_data['credit_history_length'] = all_data['credit_history_length'].apply(lambda x: x if x > 0 else np.nan) # 确保是正数

all_data = all_data.drop('earliesCreditLine', axis=1) # 删除原始日期列

# 2.2 employmentLength (就业年限) - 修复了ValueError
# 定义一个函数来清理 employmentLength
def clean_employment_length(length_str):
    if pd.isna(length_str): # 处理 NaN 值
        return np.nan
    s = str(length_str).strip().lower() # 转换为小写字符串并去除空格

    if '10+' in s:
        return 10.0
    elif '< 1' in s:
        return 0.0
    else:
        # 使用正则表达式提取数字
        match = re.search(r'(\d+)', s)
        if match:
            return float(match.group(1))
    return np.nan # 对于无法解析的字符串，返回 NaN

all_data['employmentLength'] = all_data['employmentLength'].apply(clean_employment_length)


# 2.3 FICO分数处理
all_data['fico_avg'] = (all_data['ficoRangeLow'] + all_data['ficoRangeHigh']) / 2

# 2.4 匿名特征 n0-n14
# 这些特征通常直接使用，但可以检查它们的分布和缺失值情况。
# 假设它们是数值型。

# 2.5 其他比率/组合特征
all_data['loanAmnt_to_annualIncome'] = all_data['loanAmnt'] / (all_data['annualIncome'] + 1e-6)
all_data['installment_to_annualIncome'] = all_data['installment'] / (all_data['annualIncome'] + 1e-6)
all_data['debt_to_income_ratio'] = all_data['dti'] / (all_data['annualIncome'] + 1e-6) # 债务收入比和年收入结合
all_data['revol_bal_to_annualIncome'] = all_data['revolBal'] / (all_data['annualIncome'] + 1e-6)

# 2.6 policyCode 字段
# 观察发现 policyCode 几乎都是1，可能没有区分度，可以考虑删除或者保留
# all_data = all_data.drop('policyCode', axis=1) # 示例：如果认为无用可以删除

print(f"特征工程后数据大小: {all_data.shape}")

# --- 3. 缺失值处理 ---
print("\n--- 3. 缺失值处理 ---")

# 统计缺失值
missing_counts = all_data.isnull().sum()
missing_cols = missing_counts[missing_counts > 0].index.tolist()
print(f"存在缺失值的列: {missing_cols}")

# 简单的缺失值填充策略：
# 数值型特征：中位数填充
# 类别型特征：众数填充或填充为'Missing'
for col in missing_cols:
    if all_data[col].dtype in ['int64', 'float64']:
        # 对于数值型，使用中位数填充
        all_data[col] = all_data[col].fillna(all_data[col].median())
    else:
        # 对于类别型（object），使用众数填充或特定字符串填充
        # 这里我们先填充为 'Missing'，后续进行编码
        all_data[col] = all_data[col].fillna('Missing')

print("缺失值处理完成。")
print(f"处理后缺失值总数: {all_data.isnull().sum().sum()}")


# --- 4. 类别特征编码 ---
print("\n--- 4. 类别特征编码 ---")

# 识别类别型特征
# 排除数值型特征和已经处理过的日期特征
categorical_features = all_data.select_dtypes(include=['object']).columns.tolist()

# 此外，一些数值型特征实际上是类别型，例如：
# grade, subGrade (虽然是等级，但可以视为类别，或者进行有序编码)
# regionCode, postCode (邮政编码和地区编码通常是类别)
# verificationStatus, initialListStatus, applicationType
# employmentTitle, purpose, title (这些可能高基数，需要注意)

# 手动添加一些被识别为数值但实际是类别的特征
potential_categorical = ['grade', 'subGrade', 'regionCode', 'postCode',
                         'verificationStatus', 'initialListStatus', 'applicationType',
                         'policyCode'] # policyCode 如果没被删除

for col in potential_categorical:
    if col in all_data.columns and all_data[col].dtype != 'object':
        all_data[col] = all_data[col].astype('object')
        if col not in categorical_features:
            categorical_features.append(col)

# 移除已经处理过的日期特征，如果它们被错误地识别为object
categorical_features = [f for f in categorical_features if f not in ['issueDate_year', 'issueDate_month', 'earliesCreditLine_year', 'earliesCreditLine_month']]

print(f"识别到的类别特征: {categorical_features}")

# 对类别特征进行One-Hot编码
# 对于高基数的特征 (如 employmentTitle, purpose, title, postCode)，One-Hot编码可能导致维度爆炸
# 可以考虑：
# 1. 频率编码 (Frequency Encoding)
# 2. 目标编码 (Target Encoding)
# 3. 限制One-Hot的类别数量 (例如，只对出现次数最多的N个类别进行One-Hot，其余归为'Other')
# 4. 直接使用LabelEncoder (对于树模型通常可以，但可能引入虚假序关系)

# 这里为了简化，我们先使用 pd.get_dummies 进行One-Hot编码
# 对于高基数特征，可以考虑更高级的编码方式
for col in categorical_features:
    # 限制高基数特征的One-Hot编码，例如只保留前50个最常见的类别，其余归为'Other'
    if all_data[col].nunique() > 50: # 假设超过50个唯一值就认为是高基数
        top_categories = all_data[col].value_counts().nlargest(50).index.tolist()
        all_data[col] = all_data[col].apply(lambda x: x if x in top_categories else 'Other')
        print(f"对高基数特征 '{col}' 进行了限制性One-Hot编码 (Top 50 + Other)。")

print("开始One-Hot编码...")
all_data = pd.get_dummies(all_data, columns=categorical_features, dummy_na=False)
print("One-Hot编码完成。")
print(f"编码后数据大小: {all_data.shape}")

# 分离训练集和测试集
X = all_data.iloc[:len(y_train)]
X_test = all_data.iloc[len(y_train):]

# 检查列是否完全匹配
if not X.columns.equals(X_test.columns):
    print("警告：训练集和测试集特征列不完全匹配！")
    # 可以通过以下方式对齐列
    train_cols = set(X.columns)
    test_cols = set(X_test.columns)

    missing_in_test = list(train_cols - test_cols)
    for col in missing_in_test:
        X_test[col] = 0 # 在测试集中添加缺失的列并填充0

    missing_in_train = list(test_cols - train_cols)
    for col in missing_in_train:
        X[col] = 0 # 在训练集中添加缺失的列并填充0

    # 重新对齐列顺序
    X_test = X_test[X.columns]
    print("列已对齐。")

print(f"最终训练集特征数量: {X.shape[1]}")
print(f"最终测试集特征数量: {X_test.shape[1]}")

del all_data
gc.collect()

# --- 5. 模型训练 (LightGBM) ---
print("\n--- 5. 模型训练 (LightGBM) ---")

# 划分训练集和验证集
# 使用 stratify=y 确保训练集和验证集的标签分布一致
X_train, X_val, y_train_split, y_val_split = train_test_split(X, y_train, test_size=0.2, random_state=42, stratify=y_train)

# LightGBM 参数设置
# 这些参数是常用的起始点，可以根据实际效果进行调优
lgb_params = {
    'objective': 'binary',        # 二分类问题
    'metric': 'auc',              # 评估指标为AUC
    'boosting_type': 'gbdt',      # 梯度提升树
    'n_estimators': 2000,         # 迭代次数，可以设置大一些，配合early_stopping
    'learning_rate': 0.01,        # 学习率
    'num_leaves': 31,             # 每棵树的最大叶子数
    'max_depth': -1,              # 树的最大深度，-1表示无限制
    'seed': 42,                   # 随机种子
    'n_jobs': -1,                 # 使用所有CPU核心
    'verbose': -1,                # 不打印训练过程中的信息
    'colsample_bytree': 0.7,      # 每棵树随机选择的特征比例
    'subsample': 0.7,             # 每次迭代随机选择的样本比例
    'reg_alpha': 0.1,             # L1正则化
    'reg_lambda': 0.1,            # L2正则化
    'min_child_samples': 20,      # 叶子节点最小样本数
    'feature_fraction': 0.7,      # 每次迭代随机选择的特征子集比例
}

model = lgb.LGBMClassifier(**lgb_params)

print("开始训练模型...")
model.fit(X_train, y_train_split,
          eval_set=[(X_val, y_val_split)],
          eval_metric='auc',
          callbacks=[lgb.early_stopping(100, verbose=False)]) # 100轮内验证集AUC没有提升则停止

print("模型训练完成。")

# 评估验证集上的AUC
val_preds = model.predict_proba(X_val)[:, 1]
val_auc = roc_auc_score(y_val_split, val_preds)
print(f"验证集AUC: {val_auc:.4f}")

del X_train, X_val, y_train_split, y_val_split
gc.collect()

# --- 6. 预测与提交 ---
print("\n--- 6. 预测与提交 ---")

print("开始对测试集进行预测...")
test_preds_proba = model.predict_proba(X_test)[:, 1] # 获取预测为1的概率

# 创建提交文件
submission = pd.DataFrame({'id': test_id, 'isDefault': test_preds_proba})

# 保存为csv文件
submission_filename = 'submission.csv'
submission.to_csv(submission_filename, index=False)

print(f"预测完成，提交文件已保存为 '{submission_filename}'")
print("提交文件前5行:")
print(submission.head())

print("\n--- 任务完成 ---")
